/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: dispatcher aicpu
 */

#include <memory>
#include "hccl_common.h"
#include "stream_pub.h"
#include "rts_notify.h"
#include "rtsq_interact/aicpu_hccl_sqcqv1.h"
#include "rtsq_interact/aicpu_hccl_sqcqv2.h"
#include "debug/dfx/profiling/profiling_manager.h"
#include "dispatcher_aicpu.h"
#include "sal.h"

namespace hccl {
constexpr uint64_t NANOSECOND_TO_SECOND = 1000000000U;
constexpr uint16_t TURN_LEFT_SHIFT_BIT = 16;
constexpr uint32_t HCCL_PER_LAUNCH_SQE_CNT = 128U; // 每编排N个SQE，做一次launchtask
#ifdef __cplusplus
extern "C" {
#endif  // __cplusplus
HcclResult HcclDispatcherAicpuInit(HcclDispatcher *dispatcher, const u32 devPhyId, DispatcherType type)
{
    CHK_PTR_NULL(dispatcher);
    DispatcherPub *pDispatcher = nullptr;
    if (type == DispatcherType::DISPATCHER_AICPU) {
        pDispatcher = new (std::nothrow) DispatcherAiCpu(devPhyId);
    } else {
        HCCL_ERROR("[HcclCommAicpu][HcclDispatcherInit] Not support the dispatcher type[%d]", type);
        return HCCL_E_NOT_SUPPORT;
    }
    CHK_PTR_NULL(pDispatcher);
    HcclResult ret = pDispatcher->Init();
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HcclCommAicpu][HcclDispatcherInit] Dispatcher init failed, type[%d]", type);
        delete pDispatcher;
        pDispatcher = nullptr;
        return ret;
    }
    *dispatcher = pDispatcher;
    return HCCL_SUCCESS;
}
#ifdef __cplusplus
}
#endif // __cplusplus

DispatcherAiCpu::DispatcherAiCpu(const u32 devPhyId)
    : DispatcherPub(INVALID_INT)
{
    aicpuInfo_.devId = devPhyId;
}

DispatcherAiCpu::~DispatcherAiCpu() {}

HcclResult DispatcherAiCpu::Init()
{
    if (aicpuInfo_.devType == DevType::DEV_TYPE_310P1 || aicpuInfo_.devType == DevType::DEV_TYPE_310P3) {
        addOneNotifyWaitSqe_ = AddOneNotifyWaitSqeV2;
        addOneRecordSqe_ = AddOneRecordSqeV2;
        addOneWriteValueRecordSqe_ = AddOneWriteValueRecordSqeV2;
        addOneMemcpySqe_ = AddOneMemcpySqeV2;
        addOneEventResetSqe_ = AddOneEventResetSqeV2;
        addOneEventRecordSqe_ = AddOneEventRecordSqeV2;
        addOneEventWaitSqe_ = AddOneEventWaitSqeV2;
    } else {
        addOneNotifyWaitSqe_ = AddOneNotifyWaitSqeV1;
        addOneRecordSqe_ = AddOneRecordSqeV1;
        addOneWriteValueRecordSqe_ = AddOneWriteValueRecordSqeV1;
        addOneMemcpySqe_ = AddOneMemcpySqeV1;
        addOneEventResetSqe_ = AddOneEventResetSqeV1;
        addOneEventRecordSqe_ = AddOneEventRecordSqeV1;
        addOneEventWaitSqe_ = AddOneEventWaitSqeV1;
        addOneRdmaDbSendSqe_ = AddOneRdmaDbSendSqeV1;
        addOneFlipPlaceHolderSqe_ = AddOneFlipPlaceHolderSqeV1;
        CHK_PTR_NULL(addOneRdmaDbSendSqe_);
        CHK_PTR_NULL(addOneFlipPlaceHolderSqe_);
    }

    CHK_PTR_NULL(addOneNotifyWaitSqe_);
    CHK_PTR_NULL(addOneRecordSqe_);
    CHK_PTR_NULL(addOneWriteValueRecordSqe_);
    CHK_PTR_NULL(addOneMemcpySqe_);
    CHK_PTR_NULL(addOneEventResetSqe_);
    CHK_PTR_NULL(addOneEventRecordSqe_);
    CHK_PTR_NULL(addOneEventWaitSqe_);

    CHK_RET(GetNotifyMaxWaitTime());
    return HCCL_SUCCESS;
}

u32 DispatcherAiCpu::GetMaxNotifyWaitTime()
{
    return notifyMaxWaitTime_;
}

HcclResult DispatcherAiCpu::WaitValue(hccl::Stream &stream, u64 waitAddr, u64 valueAddr, bool reset)
{
    u32 turnNum= *(reinterpret_cast<u32*>(static_cast<uintptr_t>(valueAddr)));
    HCCL_DEBUG("[DispatcherAiCpu][WaitValue] turnNum %u", turnNum);
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;
    CHK_RET(stream.GetNextSqeBufferAddr(sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    const HcclComStreamInfo *streamInfo;
    CHK_RET(stream.GetStreamInfo(streamInfo));
    AddOneWaitStartSqe(streamInfo->actualStreamId, taskId, waitAddr, valueAddr,
        reset, reinterpret_cast<rtStarsCcoreWaitStartSqe_t *>(sqeBuffer), sqeTypeAddr);
    std::shared_ptr<hccl::HcclSqeContext> sqeCtx;
    (void)stream.GetSqeContext(sqeCtx); //当前必返回SUCCESS
    if (sqeCtx.get() == nullptr) {
        HCCL_ERROR("[DispatcherAiCpu][WaitValue] AddCcoreWait sqeCtx is nullptr");
        return HCCL_E_INTERNAL;
    }
    sqeCtx.get()->buffer.addInfo[taskId % hccl::HCCL_SQE_MAX_CNT] = 
        ((turnNum << TURN_LEFT_SHIFT_BIT) + static_cast<uint32_t>(reset));
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::WriteValue(hccl::Stream &stream, u64 writeAddr, u64 valueAddr)
{
    u32 turnNum= *(reinterpret_cast<u32*>(static_cast<uintptr_t>(valueAddr)));
    HCCL_DEBUG("[DispatcherAiCpu][WriteValue] turnNum %u", turnNum);
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;
    CHK_RET(stream.GetNextSqeBufferAddr(sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    const HcclComStreamInfo *streamInfo;
    CHK_RET(stream.GetStreamInfo(streamInfo));
    AddOneWriteValueStartSqe(streamInfo->actualStreamId, taskId, writeAddr, valueAddr, reinterpret_cast<rtStarsCcoreWriteValueSqe_t *>(sqeBuffer),
        sqeTypeAddr);
    std::shared_ptr<hccl::HcclSqeContext> sqeCtx;
    (void)stream.GetSqeContext(sqeCtx); //当前必返回SUCCESS
    if (sqeCtx.get() == nullptr) {
        HCCL_ERROR("[DispatcherAiCpu][WriteValue] AddCcoreNotify sqeCtx is nullptr");
        return HCCL_E_INTERNAL;
    }
    sqeCtx.get()->buffer.addInfo[taskId % hccl::HCCL_SQE_MAX_CNT] = turnNum;
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::SignalRecord(HcclRtNotify signal, hccl::Stream &stream, u32 userRank, u64 offset, s32 stage,
    bool inchip, u64 signalAddr, u32 notifyId)
{
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;
    CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
    dfxInfo->opRingBufferIdx = opRingBufferIdx_;
    dfxInfo->remoteRank = userRank;
    dfxInfo->notifyId = notifyId;
    if (inchip) {
        addOneRecordSqe_(streamInfo.actualStreamId, taskId, notifyId, sqeBuffer, sqeTypeAddr);
    } else {
        addOneWriteValueRecordSqe_(streamInfo.actualStreamId, taskId, signalAddr, sqeBuffer, sqeTypeAddr);
    }

    HCCL_INFO("[DispatcherAiCpu][SignalRecord] streamId[%d] userRank[%u] inchip[%d] devType[%d], notifyId[%u]",
        streamInfo.actualStreamId, userRank, inchip, aicpuInfo_.devType, notifyId);
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::SignalWait(HcclRtNotify signal, Stream &stream, u32 userRank, u32 remoteUserRank, s32 stage,
    bool inchip, u32 notifyId, u32 timeOut)
{
    (void)timeOut;
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;
    CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
    dfxInfo->opRingBufferIdx = opRingBufferIdx_;
    dfxInfo->remoteRank = remoteUserRank;
    dfxInfo->notifyId = notifyId;

    if (inchip || (aicpuInfo_.devType != DevType::DEV_TYPE_310P1 && aicpuInfo_.devType != DevType::DEV_TYPE_310P3)) {
        addOneNotifyWaitSqe_(streamInfo.actualStreamId, taskId, notifyId, sqeBuffer, sqeTypeAddr, dfxTimeOutConfig_);
    } else {
        u32 notifyRevisedOffset = 15U; // eventid偏移15位后为1
        u32 notifyGetEventId = 0x3FFU; // 取低15位
        if ((notifyId >> notifyRevisedOffset) != 0) {
            addOneEventWaitSqe_(streamInfo.actualStreamId,
                (notifyId & notifyGetEventId), taskId, sqeBuffer, sqeTypeAddr);

            uint8_t *sqeBuffer1 = nullptr;
            uint8_t *sqeTypeAddr1 = nullptr;
            uint8_t *sqeDfxInfoAddr1 = nullptr;
            CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer1, sqeTypeAddr1, sqeDfxInfoAddr1, taskId));
            AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr1;
            dfxInfo->opRingBufferIdx = opRingBufferIdx_;
            dfxInfo->remoteRank = INVALID_VALUE_RANKID;

            u64 addr = 0;
            addOneEventResetSqe_(streamInfo.actualStreamId,
                (notifyId & notifyGetEventId), taskId, aicpuInfo_.devId, 0,
                addr, sqeBuffer1, sqeTypeAddr1);
        } else {
            HCCL_WARNING("[DispatcherAiCpu][SignalWait] SignalWait id is not event, please check %d", notifyId);
        }
    }

    HCCL_INFO("[DispatcherAicpu][SignalWait] streamId[%u] userRank[%u] remoteRank[%u] inchip[%d] devType[%d] "
        "notifyId[%u]", streamInfo.actualStreamId, userRank, remoteUserRank, inchip, aicpuInfo_.devType, notifyId);
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::MemcpyAsync(hccl::DeviceMem &dst, const hccl::DeviceMem &src, hccl::Stream &stream,
    u32 remoteUserRank, hccl::LinkType inLinkType)
{
    // 参数有效性检查
    if (src.size() == 0) {
        HCCL_INFO("[DispatcherAiCpu][MemcpyAsync]src memory size is 0, not need copy.");
        return HCCL_SUCCESS;
    }

    if (src == dst) {
        HCCL_INFO("[DispatcherAiCpu][MemcpyAsync]src memory and dst memory is same, not need copy.");
        return HCCL_SUCCESS;
    }

    if (dst.size() < src.size()) {
        HCCL_ERROR(
            "[DispatcherAiCpu][MemcpyAsync] "\
            "The size of dst is smaller than that of src. dst addr[%p], dst size[%llu], src addr[%p], src size[%llu]",
            dst.ptr(), dst.size(), src.ptr(), src.size());
        return HCCL_E_PTR;
    }
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();

    // 将数据按4GB切分循环处理
    uint64_t spiltLoop = 0;
    uint64_t addrOffset = 0;
    uint64_t countSplit = 0;
    uint64_t countSize = src.size();
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;
    HcclReduceOp redOp = HCCL_REDUCE_RESERVED;
    rtRecudeKind_t rtReduceOp = RK_MAP_TABLE[redOp];

    if (countSize > HCCL_SDMA_MAX_COUNT_4GB) {
        spiltLoop = (countSize % HCCL_SDMA_MAX_COUNT_4GB) ? (countSize / HCCL_SDMA_MAX_COUNT_4GB) :
                                                            ((countSize / HCCL_SDMA_MAX_COUNT_4GB) - 1);
        HCCL_INFO("[DispatcherAiCpu][MemcpyAsync] MemcpyAsync SDMA task countSize is bigger than 4GB"
            " and do segmentation splitloop[%llu]", spiltLoop);
    }
    uint8_t linkType = static_cast<uint8_t>(inLinkType);
    for (uint64_t index = 0; index <= spiltLoop; index++) {
        addrOffset = index * HCCL_SDMA_MAX_COUNT_4GB;
        countSplit = (index == spiltLoop) ? (countSize - index * HCCL_SDMA_MAX_COUNT_4GB) : (HCCL_SDMA_MAX_COUNT_4GB);
        void *srcSplit = static_cast<void *>(static_cast<char *>(const_cast<void *>(src.ptr())) + addrOffset);
        void *dstSplit = static_cast<void *>(static_cast<char *>(dst.ptr()) + addrOffset);

        CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
        AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
        dfxInfo->opRingBufferIdx = opRingBufferIdx_;
        dfxInfo->remoteRank = remoteUserRank;
        dfxInfo->notifyId = INVALID_VALUE_RANKID;
        addOneMemcpySqe_(streamInfo.actualStreamId, taskId, srcSplit, countSplit , RT_DATA_TYPE_FP32, rtReduceOp,
            dstSplit, 0, aicpuInfo_.ssid, aicpuInfo_.devId, aicpuInfo_.overflowAddr, linkType, sqeBuffer, sqeTypeAddr);

        HCCL_INFO("[DispatcherAiCpu][MemcpyAsync] Call MemcpyAsync. Split para: linkType[%u], "
            "srcSplit[%p], dstSplit[%p], countSplit [%llu], taskId[%u], streamId[%u], rtDatatType[%d], rtReduceOp[%d]",
            linkType, srcSplit, dstSplit, countSplit , taskId, streamInfo.actualStreamId, RT_DATA_TYPE_FP32, rtReduceOp);
    }

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::LaunchTask(Stream &stream, bool isBlockLaunch)
{
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();
    HcclSqeContext *sqeContext = stream.GetSqeContextPtr();
    SqeRingBuffer *sqeContextBuffer = &(sqeContext->buffer);
    const auto cnt = sqeContextBuffer->sqeCnt;
    if (cnt == 0) {
        CHK_PRT_CONT(isBlockLaunch,
            HCCL_DEBUG("no sqe, streamId:%d, sqId:%u", streamInfo.actualStreamId, streamInfo.sqId));
        return HCCL_SUCCESS;
    } else if (cnt > streamInfo.sqDepth) {
        HCCL_ERROR("LaunchTask fail, cnt:%u should be less than sqDepth:%u", cnt, streamInfo.sqDepth);
        return HCCL_E_PTR;
    }

    auto &head = sqeContextBuffer->sqHead;
    auto &tail = sqeContextBuffer->sqTail;
    u32 newTail = (tail + cnt) % streamInfo.sqDepth;
    // 仅在阻塞下发场景打印，避免非阻塞场景调用时刷屏
    CHK_PRT_CONT(isBlockLaunch,
        HCCL_INFO("Before send sqid:%d cnt:%u head:%u curtail:%u newTail:%u",
        streamInfo.sqId, cnt, head, tail, newTail));

    u64 startUsec = GetCurAicpuTimestamp();
    u64 lastUsec = startUsec;
    while (((tail < head ? streamInfo.sqDepth : 0U) + tail - head + cnt >= streamInfo.sqDepth) && (tail != head)) { // 判断剩余sqe空间是否足够下发
        // 需要放在while循环进来后第一个执行
        CHK_RET(QuerySqStatusByType(aicpuInfo_.devId, streamInfo.sqId, DRV_SQCQ_PROP_SQ_HEAD, head));

        // 非阻塞下发场景，rtsq队列空间不足时直接返回
        if (isBlockLaunch == false) {
            return HCCL_SUCCESS;
        }

        // 当前流无法下发，把其他流都launch一遍，避免等待的其他流没有launch
        for (auto it = streamMap_.begin(); it != streamMap_.end(); ++it) {
            if (it->first != streamInfo.actualStreamId) {
                CHK_RET(LaunchTask(it->second, false));
            }
        }

        u64 curUsec = GetCurAicpuTimestamp();
        if (dfxTimeOutConfig_.sqFullWaitTimeOut != 0 && 
            (curUsec - startUsec > NANOSECOND_TO_SECOND * dfxTimeOutConfig_.sqFullWaitTimeOut)) {
            HCCL_ERROR("Rtsq full, timeout %lus. curhead:%u, sqId:%d", dfxTimeOutConfig_.sqFullWaitTimeOut, head,
                streamInfo.sqId);
            return HCCL_E_AGAIN;
        }

        // 等待下发阶段，每隔30s打印一次状态
        if (curUsec - lastUsec > NANOSECOND_TO_SECOND * dfx::kPrintSqInterval) {
            lastUsec = curUsec;
            HCCL_RUN_INFO("[LaunchTask][WaitLaunchWhileLoop]Current state. sqid:%d, head:%u, tail:%u, cnt:%u",
                streamInfo.sqId, head, tail, cnt);
        }

        // 下发过程中出现cqe异常
        if (checkOpExecStatusCallback_ != nullptr) {
            HcclResult opExecStatus = checkOpExecStatusCallback_();
            CHK_PRT_RET(opExecStatus != HCCL_SUCCESS,
                HCCL_ERROR("hccl aicpu stop launch for task exception or stop command, ret:%d", opExecStatus),
                opExecStatus);
        }
    }

    uint32_t left = streamInfo.sqDepth - tail;                     // sqeAddr 剩余空间
    const auto tailSqeIdx = sqeContextBuffer->tailSqeIdx;
    HCCL_INFO("cpy sqe, left:%u, tailSqeId:%u, cnt:%u", left, tailSqeIdx, cnt);
    if (cnt <= left) { // 剩余buffer放得下新增sqe
        CHK_SAFETY_FUNC_RET(memcpy_s(
            reinterpret_cast<uint8_t *>(streamInfo.sqBaseAddr) + tail * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt) * HCCL_SQE_SIZE,
            cnt * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsMirrorBuffer + tail * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt) * HCCL_SQE_SIZE,
            cnt * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsqSqeType + tail, left,
            sqeContextBuffer->sqeType + (tailSqeIdx - cnt), cnt));
        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsDfxInfo + tail, left * sizeof(AicpuDfxInfo),
            sqeContextBuffer->dfxInfo + (tailSqeIdx - cnt), cnt * sizeof(AicpuDfxInfo)));
    } else {
        CHK_SAFETY_FUNC_RET(memcpy_s(reinterpret_cast<uint8_t *>(streamInfo.sqBaseAddr) + tail * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt) * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(reinterpret_cast<uint8_t *>(streamInfo.sqBaseAddr),
            streamInfo.sqDepth * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt + left) * HCCL_SQE_SIZE,
            (cnt - left) * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsMirrorBuffer + tail * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt) * HCCL_SQE_SIZE,
            left * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsMirrorBuffer,
            streamInfo.sqDepth * HCCL_SQE_SIZE,
            sqeContextBuffer->localBuff + (tailSqeIdx - cnt + left) * HCCL_SQE_SIZE,
            (cnt - left) * HCCL_SQE_SIZE));

        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsqSqeType + tail,
            left, sqeContextBuffer->sqeType + (tailSqeIdx - cnt), left));
        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsqSqeType + 0, streamInfo.sqDepth,
            sqeContextBuffer->sqeType + (tailSqeIdx - cnt + left), (cnt - left)));
        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsDfxInfo + tail,
            left * sizeof(AicpuDfxInfo), sqeContextBuffer->dfxInfo + (tailSqeIdx - cnt), left * sizeof(AicpuDfxInfo)));
        CHK_SAFETY_FUNC_RET(memcpy_s(sqeContextBuffer->rtsDfxInfo + 0, streamInfo.sqDepth * sizeof(AicpuDfxInfo),
            sqeContextBuffer->dfxInfo + (tailSqeIdx - cnt + left), (cnt - left) * sizeof(AicpuDfxInfo)));
    }
    CHK_RET(ConfigSqStatusByType(aicpuInfo_.devId, streamInfo.sqId, DRV_SQCQ_PROP_SQ_TAIL, newTail));
    tail = newTail;
    HCCL_INFO("After send sqe:%d, sqe_num:%u, curHead:%u, curtail:%u", streamInfo.sqId, cnt, head, tail);
    sqeContextBuffer->sqeCnt = 0;
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::LaunchTasksEx(hccl::Stream &stream, std::vector<Stream> &subStreams)
{
    /* 两阶段模式，主流待正式执行时再下 */
    /* 一阶段第一次，可以先下主流 */
    HcclResult ret = LaunchTask(stream, true);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[DispatcherAiCpu][LaunchTasksEx] "\
                   "launch task failed, sqid:%u, ret:%u", stream.sqId(), ret);
        return ret;
    }

    for (u32 index = 0; index < subStreams.size(); index++) {
        ret = LaunchTask(subStreams[index], true);
        if (ret != HCCL_SUCCESS) {
            HCCL_ERROR("[DispatcherAiCpu][LaunchTasksEx] "\
                       "launch task failed, sqid:%u, ret:%u", subStreams[index].sqId(), ret);
            return ret;
        }
    }

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::LaunchAllTasks()
{
    for (auto it = streamMap_.begin(); it != streamMap_.end(); ++it) {
        HcclResult ret = LaunchTask(it->second, true);
        if (ret != HCCL_SUCCESS) {
            HCCL_ERROR("DispatcherAiCpu][LaunchAllTasks] "\
                "launch task failed, sqid:%u, ret:%u", it->second.sqId(), ret);
            return ret;
        }
    }
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::ReduceAsync(const void *src, void *dst, u64 dataCount, const HcclDataType datatype,
    HcclReduceOp redOp, Stream &stream, HcclReduceType reduceType)
{
    return (reduceType == HcclReduceType::HCCL_INLINE_REDUCE) ?
        InlineReduceAsync(src, dataCount, datatype, redOp, stream, dst) :
        TbeReduceAsync(src, dst, dataCount, datatype, redOp, stream, dst);
}

HcclResult DispatcherAiCpu::InlineReduceAsync(const void *src, u64 dataCount, const HcclDataType datatype,
    HcclReduceOp redOp, hccl::Stream &stream, void *dst, u32 remoteUserRank, hccl::LinkType inLinkType)
{
    // 参数有效性检查
    CHK_PTR_NULL(stream.ptr());
    if (dataCount == 0) {
        HCCL_INFO("[DispatcherAiCpu][InlineReduceAsync]src memory size is 0, not need inline reduce.");
        return HCCL_SUCCESS;
    }
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();

    rtDataType_t rtDataType = DT_MAP_TABLE[datatype];
    rtRecudeKind_t rtReduceOp = RK_MAP_TABLE[redOp];

    // 将数据按4GB切分循环处理
    uint64_t spiltLoop = 0;
    uint64_t addr_offset = 0;
    uint64_t countSplit = 0;
    uint64_t countSize = dataCount * SIZE_TABLE[datatype];
    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;

    if (countSize > HCCL_SDMA_MAX_COUNT_4GB) {
        spiltLoop = (countSize % HCCL_SDMA_MAX_COUNT_4GB) ? (countSize / HCCL_SDMA_MAX_COUNT_4GB) :
                                                            ((countSize / HCCL_SDMA_MAX_COUNT_4GB) - 1);
        HCCL_INFO("[DispatcherAiCpu][InlineReduceAsync] InlineReduceAsync SDMA task countSize is bigger than 4GB"
            " and do segmentation splitloop[%llu]",
            spiltLoop);
    }
    uint8_t linkType = static_cast<uint8_t>(inLinkType);
    for (uint64_t index = 0; index <= spiltLoop; index++) {
        addr_offset = index * HCCL_SDMA_MAX_COUNT_4GB;
        countSplit = (index == spiltLoop) ? (countSize - index * HCCL_SDMA_MAX_COUNT_4GB) : (HCCL_SDMA_MAX_COUNT_4GB);
        void *srcSplit = static_cast<void *>(static_cast<char *>(const_cast<void *>(src)) + addr_offset);
        void *dstSplit = static_cast<void *>(static_cast<char *>(dst) + addr_offset);

        CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
        AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
        dfxInfo->opRingBufferIdx = opRingBufferIdx_;
        dfxInfo->remoteRank = remoteUserRank;
        dfxInfo->notifyId = INVALID_VALUE_RANKID;
        addOneMemcpySqe_(streamInfo.actualStreamId, taskId, srcSplit, countSplit, rtDataType, rtReduceOp, dstSplit, 0,
            aicpuInfo_.ssid, aicpuInfo_.devId, aicpuInfo_.overflowAddr, linkType, sqeBuffer, sqeTypeAddr);

        HCCL_INFO("[DispatcherAiCpu][InlineReduceAsync] Call InlineReduceAsync. Split para: linkType[%u] "
            "srcSplit[%p], dstSplit[%p], countSplit[%llu], taskId[%u], streamId[%u], rtDatatType[%d], rtReduceOp[%d]",
            linkType, srcSplit, dstSplit, countSplit, taskId, streamInfo.actualStreamId, rtDataType, rtReduceOp);
    }

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::TbeReduceAsync(const void *src1, const void *src2, u64 count, const HcclDataType datatype,
    HcclReduceOp redOp, Stream &stream, const void *dst)
{
    HCCL_ERROR("[DispatcherAiCpu][TbeReduceAsync] aicpu do not support the tbe reduce");
    return HCCL_E_NOT_SUPPORT;
}

HcclResult DispatcherAiCpu::RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream,
    u32 remoteUserRank)
{
    CHK_RET(RdmaSend(dbindex, dbinfo, wr, stream, RdmaType::RDMA_SEND_PAYLOAD, remoteUserRank));

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream,
    u32 userRank, u64 offset)
{
    CHK_RET(RdmaSend(dbindex, dbinfo, wr, stream, RdmaType::RDMA_SEND_NOTIFY, userRank, offset));

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::RdmaSend(u32 dbIndex, u64 dbInfo, const struct send_wr &wr, hccl::Stream &stream,
    RdmaType rdmaType, u32 userRank, u64 offset)
{
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();

    uint8_t *sqeBuffer = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    uint16_t taskId = 0U;

    CHK_RET(GetStreamSqeBufferAddr(stream, sqeBuffer, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
    dfxInfo->opRingBufferIdx = opRingBufferIdx_;
    dfxInfo->remoteRank = userRank;
    dfxInfo->notifyId = INVALID_VALUE_RANKID;

    uint32_t wrLen = (wr.buf_num >= 1) ? wr.buf_list[wr.buf_num - 1].len : 0;
    addOneRdmaDbSendSqe_(streamInfo.actualStreamId, taskId, dbInfo, dbAddr_, wrLen, static_cast<uint8_t>(rdmaType),
        sqeBuffer, sqeTypeAddr);

    HCCL_INFO("[DispatcherAiCpu][RdmaSend] Call RdmaSend. para: rankId[%u] "
        "taskId[%u], streamId[%u], rtDatatType[%d], RdmaType[%d], size[%u]",
        userRank, taskId, streamInfo.actualStreamId, RT_DATA_TYPE_FP32, rdmaType, wrLen);

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::GetStreamSqeBufferAddr(hccl::Stream &stream, uint8_t *&sqeBufferAddr, uint8_t *&sqeTypeAddr,
    uint8_t *&sqeDfxInfoAddr, uint16_t &taskId)
{
    SaveStreamInfo(stream);
    HcclSqeContext* sqeContext = stream.GetSqeContextPtr();
    if (UNLIKELY(sqeContext->buffer.sqeCnt >= HCCL_PER_LAUNCH_SQE_CNT)) {
        HCCL_INFO("GetStreamSqeBufferAddr tailSqeIdx[%u], try to launchTask", sqeContext->buffer.tailSqeIdx);
        CHK_RET(LaunchTask(stream, true));
    }
    if (UNLIKELY(sqeContext->buffer.tailSqeIdx >= HCCL_SQE_MAX_CNT)) {
        CHK_RET(LaunchTask(stream, true));
        CHK_RET(dfx::ProfilingManager::ReportTaskInfo(stream));
    }
    SqeRingBuffer *sqeContextBuffer = &(sqeContext->buffer);
    uint16_t filpNum = sqeContextBuffer->filpNum;
    uint16_t nextTaskId = sqeContextBuffer->tailSqeTaskId;
    // nextTaskId=0的时候下发PlaceHolder
    if (UNLIKELY(nextTaskId == 0  && filpNum != 0)) {
        CHK_RET(AddFlipTask(stream));
    }
    if (UNLIKELY(sqeContext->buffer.tailSqeIdx >= HCCL_SQE_MAX_CNT)) {
        CHK_RET(LaunchTask(stream, true));
        CHK_RET(dfx::ProfilingManager::ReportTaskInfo(stream));
    }
    CHK_RET(stream.GetNextSqeBufferAddr(sqeBufferAddr, sqeTypeAddr, sqeDfxInfoAddr, taskId));
    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::AddFlipTask(Stream &stream)
{
    HcclSqeContext *sqeContext = stream.GetSqeContextPtr();
    SqeRingBuffer *sqeContextBuffer = &(sqeContext->buffer);
    uint16_t filpNum = sqeContextBuffer->filpNum;
    uint16_t taskId = sqeContextBuffer->tailSqeTaskId;
    CHK_RET(dfx::ProfilingManager::ReportFilpTask(stream.id(), taskId, filpNum));

    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();
 
    uint8_t *sqeBufferAddr = nullptr;
    uint8_t *sqeTypeAddr = nullptr;
    uint8_t *sqeDfxInfoAddr = nullptr;
    CHK_RET(stream.GetNextSqeBufferAddr(sqeBufferAddr, sqeTypeAddr, sqeDfxInfoAddr, taskId));
 
    AicpuDfxInfo * const dfxInfo = (AicpuDfxInfo * const)sqeDfxInfoAddr;
    dfxInfo->opRingBufferIdx = opRingBufferIdx_;
    dfxInfo->remoteRank = INVALID_VALUE_RANKID;
    dfxInfo->notifyId = INVALID_VALUE_RANKID;
    addOneFlipPlaceHolderSqe_(streamInfo.actualStreamId, filpNum, taskId, sqeBufferAddr, sqeTypeAddr);
 
    HCCL_INFO("[DispatcherAiCpu][AddFlipTask] Call AddFlipTask. para: taskId[%u], streamId[%u], filpNum[%u]]", taskId,
        streamInfo.actualStreamId, filpNum);

    return HCCL_SUCCESS;
}

HcclResult DispatcherAiCpu::AddRetryPreamble(Stream &stream)
{
    CHK_RET(AddFlipTask(stream));
    return HCCL_SUCCESS;
}

void DispatcherAiCpu::SaveStreamInfo(hccl::Stream &stream)
{
    const HcclComStreamInfo &streamInfo = stream.GetHcclStreamInfo();
    if (streamMap_.find(streamInfo.actualStreamId) == streamMap_.end()) {
        streamMap_.insert({streamInfo.actualStreamId, stream});
        HCCL_INFO("[DispatcherAiCpu][SaveStreamInfo] stream id[%d]", streamInfo.actualStreamId);
    }
    return;
}
} // namespace hccl